import asyncio
import uuid
from app.core.database import redis_pool
class DistributedLock:
    def __init__(self, lock_key: str, expire_time: int = 10):
        """
        :param lock_key: 锁的唯一标识（如 "lock:task:123"）
        :param expire_time: 锁自动过期时间（秒），防止死锁
        """
        self.redis = redis_pool
        self.lock_key = lock_key
        self.expire_time = expire_time
        self.lock_value = str(uuid.uuid4())

    async def acquire(self, blocking: bool = True, timeout: float = 10) -> bool:
        """
        异步获取锁
        :param blocking: 是否阻塞等待
        :param timeout: 最大等待时间（秒）
        :return: 是否成功获取锁
        """
        if blocking:
            start = asyncio.get_event_loop().time()
            while True:
                if await self._try_acquire():
                    return True
                if timeout is not None and asyncio.get_event_loop().time() - start >= timeout:
                    return False
                await asyncio.sleep(0.01)  # 非阻塞等待
        else:
            return await self._try_acquire()

    async def _try_acquire(self) -> bool:
        """尝试一次获取锁"""
        # 注意：aioredis 的 set 支持 nx, px 参数（v2.x）
        result = await self.redis.set(
            self.lock_key,
            self.lock_value,
            nx=True,
            ex=self.expire_time
        )
        return result is True

    async def release(self) -> None:
        """安全释放锁"""
        lua_script = """
        if redis.call("GET", KEYS[1]) == ARGV[1] then
            return redis.call("DEL", KEYS[1])
        else
            return 0
        end
        """
        await self.redis.eval(lua_script, 1, self.lock_key, self.lock_value)

    async def __aenter__(self):
        return await self.acquire(blocking=True)

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.release()